#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
An attempt at some simple, self-contained pygame-based examples.
Example 01
In short:
One static body: a big polygon to represent the ground
One dynamic body: a rotated big polygon
And some drawing code to get you going.
kne
"""
import random

import cv2
import pygame
# Box2D.b2 maps Box2D.b2Vec2 to vec2 (and so on)
from Box2D.b2 import world, dynamicBody
import numpy as np

import simple.utils as utils
from simple.Agent import Agent

# --- constants ---
# Box2D deals with meters, but we want to display pixels,
# so define a conversion factor:
TARGET_FPS = 60
SCREEN_WIDTH = 256
SCREEN_HEIGHT = 256
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT), 0, 32)
pygame.display.set_caption('Learning')
clock = pygame.time.Clock()
world = world(gravity=(0, 0), doSleep=True)
colors = {
    dynamicBody: (255, 0, 0, 255),
}
dynamic_body = world.CreateDynamicBody(position=(10, 10), angle=0)
dynamic_body.CreatePolygonFixture(box=(2, 1), density=1, friction=0.3)
dynamic_body2 = world.CreateDynamicBody(position=(5, 5), angle=0)
dynamic_body2.CreatePolygonFixture(box=(1, 1), density=1, friction=0.3)
dynamic_body3 = world.CreateDynamicBody(position=(20, 20), angle=0)
dynamic_body3.CreatePolygonFixture(box=(1, 1), density=1, friction=0.3)
dynamic_body4 = world.CreateDynamicBody(position=(15, 15), angle=0)
dynamic_body4.CreatePolygonFixture(box=(1, 1), density=1, friction=0.3)
bodiesInTheWorld = {
    'dynamic_body': dynamic_body,
    'dynamic_body2': dynamic_body2,
    'dynamic_body3': dynamic_body3,
    'dynamic_body4': dynamic_body4
}
worldContext = {
    'PPM': 10.0,  # pixels per meter
    'TARGET_FPS': TARGET_FPS,
    'TIME_STEP': 1.0 / TARGET_FPS,
    'SCREEN_WIDTH': 256,
    'SCREEN_HEIGHT': 256,
    'world': world,
    'screen': screen,
    'clock': clock,
    'colors': colors,
    'bodiesInTheWorld': bodiesInTheWorld
}
running = True
agent = Agent((worldContext['SCREEN_HEIGHT'], worldContext['SCREEN_WIDTH']))
utils.drawTheWorld(worldContext)
targetBboxTopLeft = np.array([160, 40])
while running:
    # agent starts to work
    observedImage = cv2.imread('outputImages/temp.png')
    agent.createCurrentScene(observedImage)
    # after observation, check whether the target has reached
    agent.matchObjectBetweenScene()
    isDone = agent.checkTargetFininsed(1, targetBboxTopLeft)
    if isDone:
        print('target finished')
        break
    # agent.printScene(agent.currentScene)
    agent.learnCauseEffect()
    act = agent.getAction(1, targetBboxTopLeft)
    print('performed action is:', act)
    utils.changeTheWorldByAction(act, worldContext)
    utils.drawTheWorld(worldContext)
pygame.quit()
print('Done!')
